use candle_core::{backend::BackendDevice, CudaDevice, DeviceLocation, IndexOp, MetalDevice};
use log::error;
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use tokenizers::EncodeInput;

pub use candle_core::{DType, Device, Result, Tensor};
#[derive(Debug,Clone)]
pub struct EmbeddingModelConfig {
    pub path: PathBuf,
    pub dtype: DType,
    pub device: Device,
}

impl Default for EmbeddingModelConfig {
    fn default() -> Self {
        Self { path: ".".into(), dtype: DType::F32, device: Device::Cpu }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct SerializableEmbeddingModelConfig {
    pub path: PathBuf,
    pub dtype: SerializableDType,
    pub device: SerializableDevice,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
enum SerializableDType {
    U8,
    U32,
    I64,
    BF16,
    F16,
    F32,
    F64,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
enum SerializableDevice {
    Cpu,
    Cuda { gpu_id: usize },
    Metal { gpu_id: usize },
}

impl SerializableDType {
    pub fn from_dtype(dtype: &DType) -> Self {
        match dtype {
            DType::U8 => Self::U8,
            DType::U32 => Self::U32,
            DType::I64 => Self::I64,
            DType::BF16 => Self::BF16,
            DType::F16 => Self::F16,
            DType::F32 => Self::F32,
            DType::F64 => Self::F64,
        }
    }
    pub fn to_dtype(&self) -> DType {
        match self {
            SerializableDType::U8 => DType::U8,
            SerializableDType::U32 => DType::U32,
            SerializableDType::I64 => DType::I64,
            SerializableDType::BF16 => DType::BF16,
            SerializableDType::F16 => DType::F16,
            SerializableDType::F32 => DType::F32,
            SerializableDType::F64 => DType::F64,
        }
    }
}

impl SerializableDevice {
    pub fn from_device(device: &Device) -> Self {
        match device {
            Device::Cpu => Self::Cpu,
            Device::Cuda(l) => Self::from_location(&l.location()),
            Device::Metal(l) => Self::from_location(&l.location()),
        }
    }
    pub fn from_location(location: &DeviceLocation) -> Self {
        match location {
            DeviceLocation::Cpu => Self::Cpu,
            DeviceLocation::Cuda { gpu_id } => Self::Cuda {
                gpu_id: gpu_id.to_owned(),
            },
            DeviceLocation::Metal { gpu_id } => Self::Metal {
                gpu_id: gpu_id.to_owned(),
            },
        }
    }
    pub fn to_location(&self) -> DeviceLocation {
        match self {
            SerializableDevice::Cpu => DeviceLocation::Cpu,
            SerializableDevice::Cuda { gpu_id } => DeviceLocation::Cuda {
                gpu_id: gpu_id.to_owned(),
            },
            SerializableDevice::Metal { gpu_id } => DeviceLocation::Metal {
                gpu_id: gpu_id.to_owned(),
            },
        }
    }
    pub fn to_device(&self) -> Device {
        match self {
            SerializableDevice::Cpu => Device::Cpu,
            SerializableDevice::Cuda { gpu_id } => Device::Cuda(
                CudaDevice::new(gpu_id.to_owned())
                    .map_err(|e| {
                        error!(
                            "Failed to initialize Cuda device at {}.Error : {}",
                            gpu_id, e
                        );
                        e
                    })
                    .unwrap(),
            ),
            SerializableDevice::Metal { gpu_id } => Device::Metal(
                MetalDevice::new(gpu_id.to_owned())
                    .map_err(|e| {
                        error!(
                            "Failed to initialize Cuda device at {}.Error : {}",
                            gpu_id, e
                        );
                        e
                    })
                    .unwrap(),
            ),
        }
    }
}

impl SerializableEmbeddingModelConfig {
    pub fn from_config(config: &EmbeddingModelConfig) -> Self {
        Self {
            path: config.path.clone(),
            dtype: SerializableDType::from_dtype(&config.dtype),
            device: SerializableDevice::from_device(&config.device),
        }
    }
    pub fn to_config(&self) -> EmbeddingModelConfig {
        EmbeddingModelConfig {
            path: self.path.clone(),
            dtype: self.dtype.to_dtype(),
            device: self.device.to_device(),
        }
    }
}
impl Serialize for EmbeddingModelConfig {
    fn serialize<S>(&self, serializer: S) -> std::prelude::v1::Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        SerializableEmbeddingModelConfig::from_config(self).serialize(serializer)
    }
}

impl<'de> Deserialize<'de> for EmbeddingModelConfig {
    fn deserialize<D>(deserializer: D) -> std::prelude::v1::Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        Ok(SerializableEmbeddingModelConfig::deserialize(deserializer)?.to_config())
    }
}
pub trait EmbeddingModel {
    fn load<P>(path: P, dtype: DType, device: &Device) -> anyhow::Result<Self>
    where
        Self: Sized,
        P: AsRef<Path>;
    fn load_with_config(config: &EmbeddingModelConfig) -> anyhow::Result<Self>
    where
        Self: Sized,
    {
        Self::load(&config.path, config.dtype, &config.device)
    }
    /// 对一组句子进行向量化，并返回二维Tensor。
    fn embedding_batch_tensor<'a, E>(&self, sentences: Vec<E>) -> Result<Tensor>
    where
        E: Into<EncodeInput<'a>> + Send;
    fn embedding_single_tensor<'a, E>(&self, sentence: E) -> Result<Tensor>
    where
        E: Into<EncodeInput<'a>> + Send,
    {
        let res = self.embedding_batch_tensor(vec![sentence])?;
        let res = res.i(0)?;
        Ok(res)
    }
    /// 对一组句子进行向量化，并返回二维Vector。有可能会意外地出错。
    fn embedding_batch_vec<'a, E>(&self, sentences: Vec<E>) -> Result<Vec<Vec<f32>>>
    where
        E: Into<EncodeInput<'a>> + Send,
    {
        let encoded = self.embedding_batch_tensor(sentences)?;
        match encoded.to_vec2::<f32>() {
            Ok(vec) => return Ok(vec),
            Err(_) => {
                let dims = encoded.dims();
                if dims.len() == 1 {
                    return Ok(vec![encoded.to_vec1::<f32>()?]);
                }
                if dims.len() == 2 {
                    let dim1 = dims[1];
                    let vec = (0..dim1)
                        .map(|i| encoded.i(i).unwrap().to_vec1().unwrap())
                        .collect::<Vec<Vec<f32>>>();
                    return Ok(vec);
                }
                return Err(candle_core::Error::UnexpectedNumberOfDims {
                    expected: 2,
                    got: dims.len(),
                    shape: encoded.shape().clone(),
                });
            }
        }
    }
    fn embedding_single_vec<'a, E>(&self, sentence: E) -> Result<Vec<f32>>
    where
        E: Into<EncodeInput<'a>> + Send,
    {
        let encoded = self.embedding_single_tensor(sentence)?;
        let encoded = encoded.to_vec1::<f32>()?;
        Ok(encoded)
    }
}
